#include <iostream>
#include <boost/algorithm/string.hpp>
#include <openssl/err.h>
#include "SSLCommon.h"

//#define DEBUG

namespace ba = boost::algorithm;
using namespace kc_sf;

// ------------------------------------------------------------ //

SSLCommon::SSLCommon(int sockfd, SSL *ssl, unsigned int max_line_length)
: ssl_sockfd(sockfd),
  ssl_ssl(ssl),
  ssl_max_line_length(max_line_length)
{
}

// ------------------------------------------------------------ //

CommResult SSLCommon::sendData(const std::string &what, int timeout_s)
{
	const char *buffer = what.data();
	long len = what.size();
	long done = 0, b_sent = 0;

	fd_set write_flag;
	while (done < len)
	{
		FD_ZERO( &write_flag);
		FD_SET( ssl_sockfd, &write_flag);
		struct timeval timeout;
		timeout.tv_sec = timeout_s;
		timeout.tv_usec = 0;

		int res = select(ssl_sockfd + 1, 0, &write_flag, 0, &timeout);
		if (res == -1)
		{
			std::cerr << "select() failure" << std::endl;
			return Failure;
		}
		if (res == 0)
		{
			std::cerr << "select() timeout" << std::endl;
			return Timeout;
		}
		if (res == 1)
		{
			/* Try to write */
			b_sent = SSL_write(ssl_ssl, buffer + done, len - done);

			switch (SSL_get_error(ssl_ssl, b_sent))
			{
			/* We wrote something */
			case SSL_ERROR_NONE:

				done += b_sent;
				break;

			case SSL_ERROR_WANT_WRITE:
			case SSL_ERROR_WANT_READ:
				break;

			/* Some other error */
			default:
				std::cerr << "SSL write problem" << std::endl;
				return Failure;
			}
		}
	}

	// Wyslano
#ifdef DEBUG
	std::cout << "SSL OUT: " << what << std::endl;
#endif // DEBUG

	return Success;
}

// ------------------------------------------------------------ //

CommResult SSLCommon::recvData(std::string &to, const std::string &until_str,
		int timeout_s)
{
	to.clear();

	unsigned long total_rcvd = 0;
	fd_set read_flag;

	bool read_blocked = false;

	while (true)
	{
		FD_ZERO( &read_flag);
		FD_SET( ssl_sockfd, &read_flag);
		struct timeval timeout;
		timeout.tv_sec = timeout_s;
		timeout.tv_usec = 0;

		int res = select(ssl_sockfd + 1, &read_flag, 0, 0, &timeout);
		if (res == -1)
		{
			std::cerr << "select() failure" << std::endl;
			return Failure;
		}
		if (res == 0)
		{
			std::cerr << "select() timeout" << std::endl;
			return Timeout;
		}
		if (res == 1)
		{
			const size_t Len = ssl_max_line_length+ 1;
			char * buffer = new char[Len];

			do
			{
				int b_rcvd = SSL_read(ssl_ssl, buffer, Len);
				*(buffer + b_rcvd) = '\0';

				// Otrzymano
	#ifdef DEBUG
				std::cout << "SSL IN: " << buffer << std::endl;
	#endif // DEBUG

				std::string end = buffer;

				switch (SSL_get_error(ssl_ssl, b_rcvd))
				{
					case SSL_ERROR_NONE:

						total_rcvd += b_rcvd;
						to += buffer;

						if( ba::iends_with(to, until_str) )
						{
							return Success;
						}
						break;

					case SSL_ERROR_ZERO_RETURN:

						std::cerr << "Zero return" << std::endl;

						// End of data
						SSL_shutdown(ssl_ssl);
						delete[] buffer;
						return Success;
						break;

					case SSL_ERROR_WANT_READ:
						std::cerr << "Error want read" << std::endl;

						read_blocked = true;
						delete[] buffer;
						break;

					case SSL_ERROR_WANT_WRITE:
						std::cerr << "Error want write" << std::endl;
						delete[] buffer;
						break;

					default:
						delete[] buffer;
						std::cerr << "SSL read problem" << std::endl;
						ERR_print_errors_fp(stderr);
						std::cerr << SSL_get_error(ssl_ssl, b_rcvd)
								<< ": " << ERR_reason_error_string( SSL_get_error(ssl_ssl, b_rcvd) ) << std::endl;

						return Failure;
				}
				/* We need a check for read_blocked here because
				 SSL_pending() doesn't work properly during the
				 handshake. This check prevents a busy-wait
				 loop around SSL_read() */
			} while (SSL_pending(ssl_ssl) && !read_blocked);

			delete[] buffer;
		}
	}
	return Success;
}

// ------------------------------------------------------------ //

CommResult SSLCommon::recvData(
		std::string &to,
		int timeout_s,
		unsigned long how_many_bytes )
{
	to.clear();
	if( how_many_bytes == -1 )
		how_many_bytes = ssl_max_line_length;

	unsigned long total_rcvd = 0;
	fd_set read_flag;

	bool read_blocked = false;

	while (true)
	{
		FD_ZERO( &read_flag);
		FD_SET( ssl_sockfd, &read_flag);
		struct timeval timeout;
		timeout.tv_sec = timeout_s;
		timeout.tv_usec = 0;

		int res = select(ssl_sockfd + 1, &read_flag, 0, 0, &timeout);
		if (res == -1)
		{
			std::cerr << "select() failure" << std::endl;
			return Failure;
		}
		if (res == 0)
		{
			std::cerr << "select() timeout" << std::endl;
			return Timeout;
		}
		if (res == 1)
		{
			const size_t Len = ssl_max_line_length+ 1;
			char * buffer = new char[Len];

			do
			{
				int b_rcvd = SSL_read(ssl_ssl, buffer, Len);
				*(buffer + b_rcvd) = '\0';

				// Otrzymano
	#ifdef DEBUG
				std::cout << "SSL IN: " << buffer << std::endl;
	#endif // DEBUG

				std::string end = buffer;

				switch (SSL_get_error(ssl_ssl, b_rcvd))
				{
					case SSL_ERROR_NONE:
						total_rcvd += b_rcvd;
						to += buffer;
						if( total_rcvd >= how_many_bytes )
							return Success;
						break;

					case SSL_ERROR_ZERO_RETURN:

						std::cerr << "Zero return" << std::endl;

						// End of data
						SSL_shutdown(ssl_ssl);
						delete[] buffer;
						return Success;
						break;

					case SSL_ERROR_WANT_READ:
						std::cerr << "Error want read" << std::endl;

						read_blocked = true;
						delete[] buffer;
						break;

					case SSL_ERROR_WANT_WRITE:
						std::cerr << "Error want write" << std::endl;
						delete[] buffer;
						break;

					default:
						delete[] buffer;
						std::cerr << "SSL read problem" << std::endl;
						ERR_print_errors_fp(stderr);
						std::cerr << SSL_get_error(ssl_ssl, b_rcvd)
								<< ": " << ERR_reason_error_string( SSL_get_error(ssl_ssl, b_rcvd) ) << std::endl;

						return Failure;
				}
				/* We need a check for read_blocked here because
				 SSL_pending() doesn't work properly during the
				 handshake. This check prevents a busy-wait
				 loop around SSL_read() */
			} while (SSL_pending(ssl_ssl) && !read_blocked);

			delete[] buffer;
		}
	}

	return Success;
}
